"""
====================================================================
认证中间件模块 / Authentication Middleware Module
====================================================================
功能说明：
1. 提供WebSocket和HTTP请求的token认证
2. 从Redis获取用户权限信息（部门列表）
3. 验证token有效性
4. 提供用户上下文信息

使用示例：
    # 在WebSocket端点中
    user_context = await authenticate_websocket_token(ws)
    if not user_context:
        await ws.close(code=1008)
        return

    # 使用用户的部门权限过滤数据
    departments = user_context.departments

作者: Data Analysis Team
日期: 2025-01-06
====================================================================
"""

import logging
import json
from typing import Optional, List
from dataclasses import dataclass
from fastapi import WebSocket, HTTPException, status

# 导入Redis客户端
from client.redis_client import r2

# 配置日志记录器
logger = logging.getLogger(__name__)


@dataclass
class UserContext:
    """
    用户上下文数据类

    属性:
        token: 用户token字符串
        departments: 用户有权访问的部门ID列表
        is_authenticated: 是否已认证
    """
    token: str
    departments: List[int]
    is_authenticated: bool = True

    def has_department_access(self, department_id: int) -> bool:
        """
        检查用户是否有某个部门的访问权限

        参数:
            department_id: 部门ID

        返回:
            bool: 有权限返回True，否则返回False
        """
        # 如果部门列表包含0，表示超级管理员，有所有权限
        if 0 in self.departments:
            return True
        return department_id in self.departments

    def get_department_filter(self) -> List[int]:
        """
        获取用于数据过滤的部门ID列表

        返回:
            List[int]: 部门ID列表

        说明：
        - 如果是超级管理员（包含0），返回空列表表示不过滤
        - 否则返回实际的部门ID列表
        """
        # 超级管理员不需要过滤
        if 0 in self.departments and len(self.departments) == 1:
            return []
        # 移除0后返回实际部门列表
        return [dept for dept in self.departments if dept != 0]


async def get_token_from_websocket(ws: WebSocket) -> Optional[str]:
    """
    从WebSocket连接中提取token

    参数:
        ws: WebSocket连接对象

    返回:
        Optional[str]: token字符串，如果未找到返回None

    提取顺序：
    1. 从查询参数 ?token=xxx 中获取
    2. 从请求头 Authorization: Bearer xxx 中获取
    """
    # 方法1: 从查询参数获取
    token = ws.query_params.get("token")
    if token:
        logger.debug(f"✅ 从查询参数获取到token: {token[:10]}...")
        return token

    # 方法2: 从请求头获取
    auth_header = ws.headers.get("Authorization")
    if auth_header and auth_header.startswith("Bearer "):
        token = auth_header[7:]  # 移除 "Bearer " 前缀
        logger.debug(f"✅ 从Authorization头获取到token: {token[:10]}...")
        return token

    logger.warning("⚠️  未找到token（查询参数和Authorization头均为空）")
    return None


async def validate_token_from_redis(token: str) -> Optional[dict]:
    """
    从Redis验证token并获取权限信息

    参数:
        token: 用户token字符串

    返回:
        Optional[dict]: token对应的权限数据，如果token无效返回None

    Redis数据格式：
    {
        "depart": [dept_id1, dept_id2, ...]
    }
    """
    try:
        # 从Redis获取token对应的数据
        token_data = r2.get(token)

        if not token_data:
            logger.warning(f"⚠️  Token无效或已过期: {token[:10]}...")
            return None

        # 解析JSON数据
        data = json.loads(token_data)

        if not isinstance(data, dict):
            logger.error(f"❌ Token数据格式错误（非字典类型）: {token[:10]}...")
            return None

        logger.info(f"✅ Token验证成功: {token[:10]}..., 部门: {data.get('depart', [])}")
        return data

    except json.JSONDecodeError as e:
        logger.error(f"❌ Token数据JSON解析失败: {e}")
        return None

    except Exception as e:
        logger.error(f"❌ 验证token时发生错误: {e}")
        return None


async def authenticate_websocket_token(ws: WebSocket) -> Optional[UserContext]:
    """
    验证WebSocket连接的token并返回用户上下文

    参数:
        ws: WebSocket连接对象

    返回:
        Optional[UserContext]: 用户上下文对象，如果认证失败返回None

    工作流程：
    1. 从WebSocket提取token
    2. 从Redis验证token
    3. 提取部门权限信息
    4. 构建UserContext对象
    """
    # 步骤1: 提取token
    token = await get_token_from_websocket(ws)
    if not token:
        logger.warning("❌ WebSocket认证失败: 缺少token")
        return None

    # 步骤2: 验证token
    token_data = await validate_token_from_redis(token)
    if not token_data:
        logger.warning(f"❌ WebSocket认证失败: token无效 ({token[:10]}...)")
        return None

    # 步骤3: 提取部门信息
    departments = token_data.get("depart", [])

    # 确保部门列表是整数列表
    if not isinstance(departments, list):
        logger.error(f"❌ 部门数据格式错误（非列表）: {departments}")
        return None

    # 转换为整数列表（过滤None值）
    departments = [int(d) for d in departments if d is not None]

    # 步骤4: 构建用户上下文
    user_context = UserContext(
        token=token,
        departments=departments,
        is_authenticated=True
    )

    logger.info(
        f"✅ WebSocket认证成功: token={token[:10]}..., "
        f"部门数={len(departments)}, "
        f"超级管理员={'是' if 0 in departments else '否'}"
    )

    return user_context


async def require_authentication(ws: WebSocket) -> UserContext:
    """
    WebSocket端点的认证装饰器辅助函数

    参数:
        ws: WebSocket连接对象

    返回:
        UserContext: 用户上下文对象

    异常:
        关闭WebSocket连接（code=1008）如果认证失败

    使用示例：
        @app.websocket("/ws/endpoint")
        async def endpoint(ws: WebSocket):
            user = await require_authentication(ws)
            # 后续代码可以使用 user.departments
    """
    user_context = await authenticate_websocket_token(ws)

    if not user_context:
        # 认证失败，关闭WebSocket连接
        try:
            await ws.send_json({
                "type": "error",
                "code": "AUTHENTICATION_FAILED",
                "message": "认证失败：token无效或已过期，请重新登录"
            })
        except:
            pass  # 如果发送失败，忽略

        await ws.close(code=1008)  # 1008 = Policy Violation
        raise Exception("Authentication failed")

    return user_context


# ====================================================================
# HTTP认证工具函数（可选，暂未实现）
# ====================================================================

async def get_token_from_http_header(authorization: Optional[str] = None) -> Optional[str]:
    """
    从HTTP请求头提取token

    参数:
        authorization: Authorization请求头的值

    返回:
        Optional[str]: token字符串，如果未找到返回None
    """
    if not authorization:
        return None

    if not authorization.startswith("Bearer "):
        logger.warning("⚠️  Authorization头格式错误（应为'Bearer token'）")
        return None

    token = authorization[7:]  # 移除 "Bearer " 前缀
    return token


async def authenticate_http_request(authorization: Optional[str] = None) -> UserContext:
    """
    验证HTTP请求的token（用于REST API端点）

    参数:
        authorization: Authorization请求头的值

    返回:
        UserContext: 用户上下文对象

    异常:
        HTTPException: 如果认证失败

    使用示例：
        from fastapi import Depends, Header

        @app.get("/api/data")
        async def get_data(
            user: UserContext = Depends(
                lambda auth=Header(None, alias="Authorization"):
                    authenticate_http_request(auth)
            )
        ):
            # 使用 user.departments 过滤数据
    """
    token = await get_token_from_http_header(authorization)

    if not token:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="缺少认证token",
            headers={"WWW-Authenticate": "Bearer"},
        )

    token_data = await validate_token_from_redis(token)

    if not token_data:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="token无效或已过期",
            headers={"WWW-Authenticate": "Bearer"},
        )

    departments = token_data.get("depart", [])
    departments = [int(d) for d in departments if d is not None]

    return UserContext(
        token=token,
        departments=departments,
        is_authenticated=True
    )


# ====================================================================
# 权限检查工具函数
# ====================================================================

def check_department_access(
    user: UserContext,
    required_departments: List[int],
    require_all: bool = False
) -> bool:
    """
    检查用户是否有指定部门的访问权限

    参数:
        user: 用户上下文对象
        required_departments: 需要的部门ID列表
        require_all: 是否需要全部部门权限（True=需要全部，False=需要任一）

    返回:
        bool: 有权限返回True，否则返回False

    示例：
        # 检查用户是否有部门1或部门2的权限
        has_access = check_department_access(user, [1, 2], require_all=False)

        # 检查用户是否同时有部门1和部门2的权限
        has_access = check_department_access(user, [1, 2], require_all=True)
    """
    # 超级管理员有所有权限
    if 0 in user.departments:
        return True

    if require_all:
        # 需要所有部门权限
        return all(dept in user.departments for dept in required_departments)
    else:
        # 需要任一部门权限
        return any(dept in user.departments for dept in required_departments)


def build_department_filter(user: UserContext) -> dict:
    """
    构建Elasticsearch/MongoDB查询的部门过滤条件

    参数:
        user: 用户上下文对象

    返回:
        dict: 部门过滤查询条件

    说明：
    - 超级管理员返回空字典（不过滤）
    - 普通用户返回部门ID的terms查询条件

    使用示例：
        # Elasticsearch
        must_conditions.append(build_department_filter(user))

        # MongoDB
        query.update(build_department_filter(user))
    """
    # 超级管理员不过滤
    if 0 in user.departments and len(user.departments) == 1:
        return {}

    # 移除超级管理员标识(0)，返回实际部门列表
    departments = [dept for dept in user.departments if dept != 0]

    if not departments:
        return {}

    # 返回Elasticsearch terms查询格式
    return {"terms": {"dataAuthority.depart": departments}}
